#include "StdAfx.h"
#include "CorelSymbolEngine.h"
#include <psapi.h>

extern VOID (WINAPI *TrueOutputDebugStringA)(LPCSTR lpOutputString);
extern VOID (WINAPI *TrueOutputDebugStringW)(LPCWSTR lpOutputString);
extern int (WINAPI *TrueMultiByteToWideChar)(UINT CodePage, DWORD dwFlags, LPCSTR lpMultiByteStr, int cbMultiByte, LPWSTR lpWideCharStr, int cchWideChar);
extern int (WINAPI *TrueWideCharToMultiByte)(UINT CodePage, DWORD dwFlags, LPCWSTR lpWideCharStr, int cchWideChar, LPSTR lpMultiByteStr, int cbMultiByte, LPCSTR lpDefaultChar, LPBOOL lpUsedDefaultChar);


CCorelSymbolEngine::CCorelSymbolEngine()
: m_bInited(FALSE)
{
    InitializeCriticalSection(&m_csLock);
    m_mapLoadedSymbol.clear();
}

CCorelSymbolEngine::~CCorelSymbolEngine()
{
    DeleteCriticalSection(&m_csLock);

    if (m_bInited)
    {
        m_kCorelDbgHelp.SymCleanup(GetCurrentProcess);
    }
}

BOOL CALLBACK EnumSymProc( 
                          PSYMBOL_INFO pSymInfo,   
                          ULONG SymbolSize,      
                          PVOID UserContext)
{
    char t_szTmp[1024];
    sprintf_s(t_szTmp, sizeof(t_szTmp), "%08I64X %4u %s\n", pSymInfo->Address, SymbolSize, pSymInfo->Name);
    TrueOutputDebugStringA(t_szTmp);
    return TRUE;
}


BOOL CCorelSymbolEngine::Init(BOOL bForceInit)
{
    EnterCriticalSection(&m_csLock);

    if ((FALSE == bForceInit) && (TRUE == m_bInited)) 
    {
        goto _exit;
    }

    m_mapLoadedSymbol.clear();

#ifdef _DEBUG
    m_kCorelDbgHelp.SymSetOptions(SYMOPT_LOAD_LINES|SYMOPT_DEFERRED_LOADS|SYMOPT_UNDNAME|SYMOPT_DEBUG);
#else
    m_kCorelDbgHelp.SymSetOptions(SYMOPT_LOAD_LINES|SYMOPT_UNDNAME|SYMOPT_DEFERRED_LOADS);
#endif
    
    m_kCorelDbgHelp.SymInitializeW(GetCurrentProcess(), SymbolSearchPath().c_str(), FALSE);

    HANDLE hSnap = CreateToolhelp32Snapshot(8, GetCurrentProcessId());
    MODULEENTRY32 me;
    me.dwSize = sizeof(MODULEENTRY32);

    BOOL bContinue = Module32First(hSnap, &me);
    while (bContinue)
    {
        char szExePath[300];
        char szModule[300];
        TrueWideCharToMultiByte(CP_ACP, WC_COMPOSITECHECK, me.szExePath, lstrlen(me.szExePath) + 1, szExePath, sizeof(szExePath), NULL, NULL);
        TrueWideCharToMultiByte(CP_ACP, WC_COMPOSITECHECK, me.szModule, lstrlen(me.szModule) + 1, szModule, sizeof(szModule), NULL, NULL);

        for (char *p = szModule ; *p != NULL ; ++p) *p = toupper(*p);
        m_mapLoadedSymbol[std::string(szModule)] = TRUE;

        SetLastError(S_OK);
        if (0 == m_kCorelDbgHelp.SymLoadModule64(GetCurrentProcess(), 0, szExePath, szModule, (DWORD64)me.modBaseAddr, me.modBaseSize) && (ERROR_SUCCESS != GetLastError()))
        {
            //sprintf(t_szTmpBuf, "[CorelSpy]Error (%d) for loading symbol(%s)\n", GetLastError(), szModule);
            //TrueOutputDebugStringA(t_szTmpBuf);
        }
        else
        {
            //sprintf(t_szTmpBuf, "[CorelSpy] Loading symbol(%s) success![Base = 0x%08X, size = %d]\n", szModule, me.modBaseAddr, me.modBaseSize);
            //TrueOutputDebugStringA(t_szTmpBuf);
        }

        bContinue = Module32Next(hSnap, &me);
    }

    m_bInited = TRUE;
    CloseHandle(hSnap);
_exit:
    
    LeaveCriticalSection(&m_csLock);

    return m_bInited;
}

BOOL CCorelSymbolEngine::LoadSymbol(HMODULE hModule, LPCSTR szModule)
{
    BOOL t_ret = TRUE;
    if (FALSE == m_bInited)
    {
        goto _exit;
    }

    if (FALSE == TryEnterCriticalSection(&m_csLock))
    {
        // ToDO : LoadSymbol next time

        goto _exit;
    }

    EnterCriticalSection(&m_csLock);

    const char *p;
    char t_pszModule[1024];
    for (p = szModule + strlen(szModule) - 1 ; 
         *p != '/' && *p != '\\' && p > szModule ; 
         --p);


    if (*p == '/' || *p == '\\') ++p;
    strcpy_s(t_pszModule, sizeof(t_pszModule), p);
    for (char *pp = t_pszModule ; *pp != NULL ; ++pp) *pp = toupper(*pp);

    if (!m_mapLoadedSymbol[std::string(t_pszModule)])
    {
        MODULEINFO t_mi;

        GetModuleInformation(GetCurrentProcess(), hModule, &t_mi, sizeof(MODULEINFO));
        SetLastError(S_OK);

        if (0 == m_kCorelDbgHelp.SymLoadModule64(GetCurrentProcess(), 0, szModule, t_pszModule, (DWORD64)t_mi.lpBaseOfDll, t_mi.SizeOfImage))
        {
            //sprintf(t_szTmpBuf, "[CorelSpy]Error (%d) for loading symbol(%s)\n", GetLastError(), t_pszModule);
            //TrueOutputDebugStringA(t_szTmpBuf);
        }
        else
        {
            //sprintf(t_szTmpBuf, "[CorelSpy] Loading symbol(%s) success![Base = 0x%08X, size = %d]\n", szModule, t_mi.lpBaseOfDll, t_mi.SizeOfImage);
            //TrueOutputDebugStringA(t_szTmpBuf);
        }

        m_mapLoadedSymbol[std::string(t_pszModule)] = TRUE;
    }

    LeaveCriticalSection(&m_csLock);

_exit:

    
    return t_ret;
}

void CCorelSymbolEngine::StackTrace(
                               PCONTEXT pContext, 
                               size_t uNumOfFrameToSkip, 
                               std::wstring &wstrStackTrace)
{
    std::string strAddress;
    std::vector<std::string> vecStackTrace;
    STACKFRAME64 stStackFrame = {0};
    DWORD dwError;

    Init(FALSE);

    EnterCriticalSection(&m_csLock);

    stStackFrame.AddrPC.Offset = pContext->Eip;
    stStackFrame.AddrPC.Mode = AddrModeFlat;
    stStackFrame.AddrFrame.Offset = pContext->Ebp;
    stStackFrame.AddrFrame.Mode = AddrModeFlat;
    stStackFrame.AddrStack.Offset = pContext->Esp;
    stStackFrame.AddrStack.Mode = AddrModeFlat;

    while (true)
    {
        BOOL bSuccess = m_kCorelDbgHelp.StackWalk64(
            IMAGE_FILE_MACHINE_I386,
            GetCurrentProcess(),
            GetCurrentThread(),
            &stStackFrame,
            pContext,
            NULL,
            (PFUNCTION_TABLE_ACCESS_ROUTINE64)m_kCorelDbgHelp.SymFunctionTableAccess64,
            (PGET_MODULE_BASE_ROUTINE64)m_kCorelDbgHelp.SymGetModuleBase64, 
            NULL);

        if (FALSE == bSuccess)
        {
            dwError = GetLastError();
            break;
        }

        if (!AddressToString((PVOID)stStackFrame.AddrPC.Offset, strAddress)) break;
        vecStackTrace.push_back(strAddress);
        strAddress.clear();
    }

    for (size_t i = uNumOfFrameToSkip ; i < vecStackTrace.size() ; ++i)
    {
        DWORD cchTarget;
        cchTarget = MultiByteToWideChar(CP_ACP, MB_PRECOMPOSED, vecStackTrace[i].c_str(), vecStackTrace[i].length() + 1, NULL, 0);
        LPWSTR lpwszStackTrace;
        lpwszStackTrace = (LPWSTR)HeapAlloc(GetProcessHeap(), 0, cchTarget * sizeof(WCHAR));
        if (NULL != lpwszStackTrace)
        {
            MultiByteToWideChar(CP_ACP, MB_PRECOMPOSED, vecStackTrace[i].c_str(), vecStackTrace[i].length() + 1, lpwszStackTrace, cchTarget);
            wstrStackTrace.append(L"\t");
            wstrStackTrace.append(lpwszStackTrace);
            wstrStackTrace.append(L"\n");
            HeapFree(GetProcessHeap(), 0, lpwszStackTrace);
            lpwszStackTrace = NULL;
        }
    }

    LeaveCriticalSection(&m_csLock);
}

std::wstring CCorelSymbolEngine::SymbolSearchPath()
{
    std::wstring t_wstrSymbolPath;
    WCHAR t_wszServerPath[MAX_PATH];
    WCHAR t_wszExePath[MAX_PATH];
    WCHAR t_wszProfilePath[MAX_PATH];
    WCHAR t_wszSymbolPath[1024];

    // Retrieve exe path
    GetModuleFileName(NULL, t_wszExePath, MAX_PATH);
    *(WCHAR*)wcsrchr(t_wszExePath, L'\\') = NULL;
    t_wstrSymbolPath.append(t_wszExePath);

    // Retrieve server path
    DWORD t_dwRet = GetEnvironmentVariableW(ENVSTR_SERVER_PATH, t_wszServerPath, MAX_PATH);
    if (t_dwRet && MAX_PATH >= t_dwRet)
    {
        t_wstrSymbolPath.append(L";");
        t_wstrSymbolPath.append(t_wszServerPath);

        // Add search path defined in ini file
        wcscpy_s(t_wszProfilePath, sizeof(t_wszProfilePath)/sizeof(WCHAR), t_wszServerPath);
        wcscat_s(t_wszProfilePath, sizeof(t_wszProfilePath)/sizeof(WCHAR), L"\\CorelSpyConsole.ini");
        GetPrivateProfileStringW(
            L"SYMBOL", 
            L"PATH", 
            L"SRV*D:\\symbols\\*http://msdl.microsoft.com/download/symbols",
            t_wszSymbolPath,
            1024,
            t_wszProfilePath);

        t_wstrSymbolPath.append(L";");
        t_wstrSymbolPath.append(t_wszSymbolPath);
    }

    return t_wstrSymbolPath;
}

BOOL CCorelSymbolEngine::AddressToString(
                                    PVOID address, 
                                    std::string &strAddress)
{
    const unsigned short cchMaxMessageLen = 1024;
    char szBuf[cchMaxMessageLen];
    BOOL t_bSuccess = FALSE;

    ULONG64 buffer[(sizeof(SYMBOL_INFO) +
        MAX_SYM_NAME * sizeof(TCHAR) +
        sizeof(ULONG64) - 1) /
        sizeof(ULONG64)];

    PSYMBOL_INFO pSymbol = (PSYMBOL_INFO) buffer;
    pSymbol->SizeOfStruct = sizeof(SYMBOL_INFO);
    pSymbol->MaxNameLen = MAX_SYM_NAME;


    HANDLE hProcess = GetCurrentProcess();
    strAddress.clear();
    sprintf_s(szBuf, sizeof(szBuf), "0x%08x : ", (DWORD)address);
    strAddress.append(szBuf);

    // module
    IMAGEHLP_MODULE64 stModuleInfo;
    stModuleInfo.SizeOfStruct = sizeof(IMAGEHLP_MODULE64);
    if (TRUE == m_kCorelDbgHelp.SymGetModuleInfo64(hProcess, (DWORD64)address, &stModuleInfo))
    {
        //strAddress.append(stModuleInfo.ImageName);
        strAddress.append(stModuleInfo.ModuleName);
        strAddress.append("!");
    }
    else
    {
        //TrueOutputDebugStringA("[CorelSpy] Cannot get module info");
        goto _exit;
    }

    // function
    szBuf[0] = NULL;
    DWORD64 dw64Displacement;
    if (m_kCorelDbgHelp.SymFromAddr(GetCurrentProcess(), (DWORD)address, &dw64Displacement, pSymbol))
    {
        strAddress.append(pSymbol->Name, pSymbol->NameLen);
        if (0 != dw64Displacement)
        {
            sprintf_s(szBuf, sizeof(szBuf), "+0x%04x", dw64Displacement);
            strAddress.append(szBuf);
            szBuf[0] = NULL;
        }
        t_bSuccess = TRUE;
    }
    else
    {
        ReportError(L"SymFromAddr Fail.", GetLastError());
        goto _exit;
    }

    // line
    IMAGEHLP_LINE stLineInfo;
    stLineInfo.SizeOfStruct = sizeof(IMAGEHLP_LINE);
    DWORD dwDisplacement = 0;
    if (m_kCorelDbgHelp.SymGetLineFromAddr(GetCurrentProcess(), (DWORD)address, &dwDisplacement, &stLineInfo))
    {
        if (0 == dwDisplacement)
        {
            sprintf_s(szBuf, sizeof(szBuf), ", %s(%d)", stLineInfo.FileName, stLineInfo.LineNumber);
        }
        else
        {
            sprintf_s(szBuf, sizeof(szBuf), ", %s(%d)+%04d byte(s)", stLineInfo.FileName, stLineInfo.LineNumber, dwDisplacement);
        }
        strAddress.append(szBuf);

#ifdef _DEBUG
        WCHAR wszDebugMsg[500];
        swprintf_s(wszDebugMsg, sizeof(wszDebugMsg)/sizeof(wszDebugMsg[0]), L"SymGetLineFromAddr, %x --> %s, %d, dis = %d ", address, stLineInfo.FileName, stLineInfo.LineNumber, dwDisplacement);
        ReportError(wszDebugMsg, 0);
#endif
    }
    else
    {
        WCHAR wszDebugMsg[1000];
        swprintf_s(wszDebugMsg, sizeof(wszDebugMsg)/sizeof(wszDebugMsg[0]), L"SymGetLineFromAddr (%x) Fail. ", address);
        ReportError(wszDebugMsg, GetLastError());
    }

_exit:;

    return t_bSuccess;
}

void CCorelSymbolEngine::ReportError(
                                LPCWSTR lpwszErrorMsg, 
                                DWORD dwErrorCode)
{
    const unsigned short cchErrorCodeLen = 1024;
    WCHAR wszErrorCode[cchErrorCodeLen];
    std::wstring wstrErrorMsg;

    swprintf_s(wszErrorCode, sizeof(wszErrorCode)/sizeof(wszErrorCode[0]), L"%d", dwErrorCode);
    wstrErrorMsg = lpwszErrorMsg;
    wstrErrorMsg.append(L", err code:");
    wstrErrorMsg.append(wszErrorCode);
}
